#!/usr/bin/env python

# 将得到的fasta文件 加上 emapper 注释

import argparse
import sys
from Bio import SeqIO

parser = argparse.ArgumentParser(
    description='''从emapper 结果 对fasta 进行 添加注释信息
    用法:
    add_emapper2fasta.py -i Cves_OGS_v3.1.cds.fasta -e emapper_out.emapper.annotations2 -o Cves_OGS_v3.2.cds.fasta''')

parser.add_argument('-i',
                help='fasta文件')

parser.add_argument('-e',
                help='emapper结果')

parser.add_argument('-o',
                help='输出文件')


args = parser.parse_args()

if not args.i or not args.e or not args.o:
    parser.print_help()
    sys.exit()



# 对gff3进行注释 给予 emapper的信息

fasta_file = args.i

emapper_file = args.e

outfile = args.o



emapper_dic = {}

head = []

with open(emapper_file) as fila:
    for i in fila:
        k = i.split('#')[0].strip().split('\t')
        if len(k)>1:
            gene_name = k[0].split('.mRNA')[0]
            target_name = k[8]
            des_name = k[7]
            emapper_dic[gene_name] = [target_name,des_name]

new_dic = {}

for i in emapper_dic:
    if emapper_dic[i][0] == '-':
        pass
    else:
        new_dic[i] = emapper_dic[i]


used_dic = {}
second_dic = {}
for i in new_dic:
    new_name = new_dic[i][0]
    if new_name in used_dic:
        used_dic[new_name]+=1
        second_dic[new_name] = 1
    else:
        used_dic[new_name] = 1
        



name_change_dic = {}

for i in new_dic:
    new_name = new_dic[i][0]
    if used_dic[new_name]>1:
        name_change_dic[i] = new_name.replace('-','')+'-'+str(second_dic[new_name])
        second_dic[new_name] += 1
    else:
        name_change_dic[i] = new_name.replace('-','')


# 修改 信息

outfile = open(outfile,'w')

for i in SeqIO.parse(fasta_file,'fasta'):
    gene_name = i.name.split('.mRNA')[0]
    i.name = i.name.split('_')[0]
    if gene_name in name_change_dic:
        d = str(i.description)
        d += (' ' + name_change_dic[gene_name])
        i.description = d

    else:
        d = str(i.description)
        d += (' ' + gene_name)
        i.description = d
    
    i.description += ' '

    if gene_name in emapper_dic and emapper_dic[gene_name][1]!='-' and len(emapper_dic[gene_name][1])<60:


        i.description += ' '+emapper_dic[gene_name][1]
        #print(i.description)
    i.description = str(i.description).strip(' ')

    SeqIO.write(i,outfile,'fasta')


outfile.close()